#include "vedb.h"

void VEDBSetup(nlohmann::json& vedbJson) {
    using namespace nlohmann;
    using namespace std;
    json& fileJson = vedbJson["files"];
    json& W = vedbJson["W"];
    json& ownerJson = vedbJson["owner"];

    gmp_randclass gmp_rand(gmp_randinit_default);
    gmp_rand.seed(std::time(nullptr));
    mpz_class KI(gmp_rand.get_z_bits(256));
    ownerJson["KI"] = KI.get_str();
    mpz_class KS(gmp_rand.get_z_bits(256));
    ownerJson["KS"] = KS.get_str();
    mpz_class KX(gmp_rand.get_z_bits(256));
    ownerJson["KX"] = KX.get_str();
    mpz_class KZ(gmp_rand.get_z_bits(256));
    ownerJson["KZ"] = KZ.get_str();

    pair_class pairing(ownerJson["param"].get<string>());
    ele_class g(pairing, "G1");
    g.random();
    ownerJson["g"] = g.to_string();
    ele_class s(pairing, "Zr");
    s.random();
    ownerJson["s"] = s.to_string();
    ele_class g1(pairing, "Zr");
    g1.random();
    ownerJson["g1"] = g1.to_string();
    ele_class g2(pairing, "Zr");
    g2.random();
    ownerJson["g2"] = g2.to_string();
    ele_class g3(pairing, "Zr");
    g3.random();
    ownerJson["g3"] = g3.to_string();

    json& TSet = ownerJson["TSet"];
    json& XSet = ownerJson["XSet"];
    json& Stag = ownerJson["Stag"];
    TSet.clear();
    XSet.clear();
    Stag.clear();
    vector<ele_class> XSet_vec;
    vector<ele_class> Stag_vec;
    BMapAcc acc(g, s);
    for (auto& w : W.items()) {
        vector<ele_class> ew;
        int c = 1;
        mpz_class Ke(F(KI.get_str(), w.key()), 16);
        ele_class num_1(pairing, "Zr");
        num_1 = 1;
        ele_class w_ele(pairing, "Zr");
        w_ele = mpz_class(H(w.key()), 16);
        string stagw = F(KS.get_str(), (g1 ^ (w_ele.inverse())).to_string());
        string h_stagw = H(stagw);
        Stag.push_back(h_stagw);
        ele_class stag_ele(pairing, "Zr");
        stag_ele.set_str(h_stagw, 16);
        Stag_vec.push_back(stag_ele);
        string l;
        for (auto& ind_js: w.value()) {
            int ind = ind_js.get<int>();
            ele_class xind(pairing, "Zr");
            xind = mpz_class(Fp(KI.get_str(), to_string(ind)), 16);
            ele_class z(pairing, "Zr");
            z = mpz_class(Fp(KZ.get_str(), (g2 ^ (w_ele.inverse())).to_string() + to_string(c)), 16);
            l = F(stagw, to_string(c));

            string e = Enc(Ke, to_string(ind));
            ele_class y(pairing, "Zr");
            y = xind * (z.inverse());
            vector<string> e_y;
            e_y.push_back(e);
            e_y.push_back(y.to_string());
            TSet[l] = e_y;
            ele_class e1(pairing, "Zr");
            e1.set_str(e, 16);
            ew.push_back(e1);
            ele_class xtagw(pairing, "G1");
            ele_class t(pairing, "Zr");
            t.set_str(F(KX.get_str(), (g3 ^ (w_ele.inverse())).to_string()), 16);
            xtagw = g ^ (t * xind);
            XSet_vec.push_back(xtagw);
            XSet.push_back(xtagw.to_string());
            ++c;
        }
        l = F(stagw, to_string(0));
        ele_class acc_ew = acc(ew);
        vector<string> t;
        t.push_back(acc_ew.to_string());
        t.push_back(H(Ke.get_str() + acc_ew.to_string()));
        TSet[l] = t;
    }
    ownerJson["acc_XSet"] = acc(XSet_vec).to_string();
    ownerJson["acc_Stag"] = acc(Stag_vec).to_string();
}

void setSearchWords(nlohmann::json& vedbJson, int num) {
    assert(num > 0);
    using namespace nlohmann;
    using namespace std;
    json& WJson = vedbJson["W"];
    json& WUserJson = vedbJson["user"]["W"];
    WUserJson.clear();
    if (num > WJson.size()) num = WJson.size() - 1;
    srand(int(time(0)));
    int begin = rand() % (WJson.size() - num);
    for (auto it = WJson.begin(); it != WJson.end(); ++it, --begin) {
        if (begin <= 0) {
            WUserJson.push_back(it.key());
            if ((--num) <= 0) {
                break;
            }
        }
    }
}

void TokenGen(nlohmann::json& vedbJson) {
    using namespace nlohmann;
    using namespace std;
    json& W = vedbJson["W"];
    json& ownerJson = vedbJson["owner"];
    json& userJson = vedbJson["user"];
    string KZ_str = ownerJson["KZ"].get<string>();
    string KS_str = ownerJson["KS"].get<string>();
    string KX_str = ownerJson["KX"].get<string>();
    pair_class pairing(ownerJson["param"].get<string>());
    ele_class g(pairing, "G1");
    g.set_str(ownerJson["g"].get<string>());
    ele_class g1(pairing, "Zr");
    g1.set_str(ownerJson["g1"].get<string>());
    ele_class g2(pairing, "Zr");
    g2.set_str(ownerJson["g2"].get<string>());
    ele_class g3(pairing, "Zr");
    g3.set_str(ownerJson["g3"].get<string>());
    ele_class w1_ele(pairing, "Zr");
    w1_ele = mpz_class(H(userJson["W"][0].get<string>()), 16);
    string stag = F(KS_str, (g1 ^ (w1_ele.inverse())).to_string());
    userJson["Td"]["stag"] = stag;

    int c = 1;
    int d = userJson["W"].size();
    userJson["Td"]["xtoken"].clear();
    while (1) {
        string l = F(stag, to_string(c));
        if (!ownerJson["TSet"].contains(l)) {
            break;
        }
        vector<string> xtokenc;
        for (int i = 1; i < d; ++i) {
            ele_class z(pairing, "Zr");
            z = mpz_class(Fp(KZ_str, (g2 ^ (w1_ele.inverse())).to_string() + to_string(c)), 16);
            ele_class t(pairing, "Zr");
            ele_class wi_ele(pairing, "Zr");
            wi_ele = mpz_class(H(userJson["W"][i].get<string>()), 16);
            t.set_str(F(KX_str, (g3 ^ (wi_ele.inverse())).to_string()), 16);
            xtokenc.push_back((g ^ (z * t)).to_string());
        }
        userJson["Td"]["xtoken"].push_back(json(xtokenc));
        ++c;
    }
}

void Search(nlohmann::json& vedbJson) {
    using namespace nlohmann;
    using namespace std;
    json& W = vedbJson["W"];
    json& ownerJson = vedbJson["owner"];
    json& userJson = vedbJson["user"];

    string stag_w1_str = userJson["Td"]["stag"].get<string>();
    string l = F(stag_w1_str, "0");
    string acc_ew = ownerJson["TSet"][l][0].get<string>();
    string hash_acc_ew = ownerJson["TSet"][l][1].get<string>();
    vector<string> Rw1;
    vector<string> y;
    int c = 1;
    while (1) {
        l = F(stag_w1_str, to_string(c));
        if (!ownerJson["TSet"].contains(l)) {
            break;
        }
        Rw1.push_back(ownerJson["TSet"][l][0].get<string>());
        y.push_back(ownerJson["TSet"][l][1].get<string>());
        ++c;
    }
    pair_class pairing(ownerJson["param"].get<string>());
    ele_class g(pairing, "G1");
    g.set_str(ownerJson["g"].get<string>());
    ele_class s(pairing, "Zr");
    s.set_str(ownerJson["s"].get<string>());
    ele_class stagw1(pairing, "Zr");
    stagw1.set_str(userJson["Td"]["stag"].get<string>());
    BMapAcc acc(g, s);
    // userJson["proof"].clear();
    userJson["proof"]["W_Rw1_u"].clear();
    userJson["proof"]["W_Rw1_w"].clear();
    userJson["proof"]["W_Rw1"].clear();
    userJson["proof"]["WB"].clear();
    if (Rw1.empty()) {
        ele_class uy(pairing, "Zr");
        uy = 1;
        ele_class f_stag_s(pairing, "Zr");
        f_stag_s = 1;
        ele_class wy(pairing, "G1");
        for (auto it = ownerJson["Stag"].begin(); it != ownerJson["Stag"].end(); ++it) {
            ele_class x(pairing, "Zr");
            x.set_str(it->get<string>(), 16);
            uy = uy * (x - stagw1);
            f_stag_s = f_stag_s * (x + s);
        }
        wy = g ^ ((f_stag_s - uy) / (stagw1 + s));
        userJson["proof"]["W_Rw1_u"] = uy.to_string();
        userJson["proof"]["W_Rw1_w"] = wy.to_string();
        return;
    } else {
        unordered_set<string> Rw1_set;
        for (auto& item : Rw1) {
            Rw1_set.insert(item);
        }
        ele_class w_Rw1(pairing, "G1");
        ele_class t(pairing, "Zr");
        for (auto it = ownerJson["Stag"].begin(); it != ownerJson["Stag"].end(); ++it) {
            string curr_stag = it->get<string>();
            if (Rw1_set.find(curr_stag) != Rw1_set.end()) {
                continue;
            }
            ele_class x(pairing, "Zr");
            x.set_str(curr_stag, 16);
            t = 1;
            t = t * (x + s);
        }
        w_Rw1 = g ^ t;
        userJson["proof"]["W_Rw1"] = w_Rw1.to_string();
        vector<string> R;
        vector<ele_class> B;
        unordered_set<string> B_set;
        auto xtoken_size = userJson["Td"]["xtoken"].size();
        for (size_t c = 0; c < xtoken_size; ++c) {
            auto& xtoken_json = userJson["Td"]["xtoken"][c];
            for (auto& xtokenc_json: xtoken_json) {
                ele_class b(pairing, "G1");
                b.set_str(xtokenc_json.get<string>());
                ele_class yc(pairing, "Zr");
                yc.set_str(y[c]);
                b = b ^ yc;
                if (userJson["XSet"].find(b.to_string()) != userJson["XSet"].end()) {
                    R.push_back(Rw1[c]);
                    B.push_back(b);
                    B_set.insert(xtokenc_json.get<string>());
                }
            }
        }
        
        if (!R.empty()) {
            ele_class w_B(pairing, "G1");
            ele_class t(pairing, "Zr");
            for (auto it = ownerJson["XSet"].begin(); it != ownerJson["XSet"].end(); ++it) {
                string curr_xset = it->get<string>();
                if (B_set.find(curr_xset) != B_set.end()) {
                    continue;
                }
                ele_class x(pairing, "G1");
                x.set_str(curr_xset);
                t = 1;
                t = t * (x + s);    // 这里有问题
            }
            w_B = g ^ t;  
            userJson["proof"]["WB"] = w_B.to_string();
        }
        userJson["R"] = json(R);
        userJson["Rw1"] = json(Rw1);
    }
}

void Verify(nlohmann::json& vedbJson) {
    using namespace nlohmann;
    using namespace std;
    json& ownerJson = vedbJson["owner"];
    json& userJson = vedbJson["user"];
    int d = userJson["W"].size();
    string KI_str = ownerJson["KI"].get<string>();
    string KX_str = ownerJson["KX"].get<string>();
    mpz_class KI(KI_str);
    mpz_class Ke(F(KI_str, userJson["W"][0].get<string>()), 16);
    userJson["xtag"].clear();
    pair_class pairing(ownerJson["param"].get<string>());
    ele_class g(pairing, "G1");
    g.set_str(ownerJson["g"].get<string>());
    ele_class s(pairing, "Zr");
    s.set_str(ownerJson["s"].get<string>());
    ele_class g3(pairing, "Zr");
    g3.set_str(ownerJson["g3"].get<string>());
cout << 1 << endl;
    for (auto& e_ind : userJson["Rw1"]) {
        string ind_str = Dec(Ke, e_ind.get<string>());
        for (auto it = ind_str.begin(); it != ind_str.end(); ++it) {
            if ((*it) < '0' || (*it) > '9') {
                ind_str.erase(it, ind_str.end());
                break;
            }
        }
        vector<string> xtagi;
        // cout << ind_str << endl;
        ele_class xind(pairing, "Zr");
        xind = mpz_class(Fp(KI_str, ind_str), 16);
        
        for (int i = 1; i < d; ++i) {
            ele_class t(pairing, "Zr");
            ele_class wi_ele(pairing, "Zr");
            wi_ele = mpz_class(H(userJson["W"][i].get<string>()), 16);
            t.set_str(F(KX_str, (g3 ^ (wi_ele.inverse())).to_string()), 16);
            xtagi.push_back((g ^ (xind * t)).to_string());
        }
        userJson["xtag"].push_back(json(xtagi));
    }
cout << 12 << endl;
    // 这里是混的 懒得写了
    ownerJson["proof"].clear();
    for (auto& xtagi : userJson["xtag"]) {
        for (int j = 0; j < d - 1; ++j) {
            if (!ownerJson["XSet"].contains(xtagi[j])) {
                ele_class uy(pairing, "Zr");
                uy = 1;
                ele_class fs(pairing, "Zr");
                fs = 1;
                ele_class wy(pairing, "G1");
                ele_class x(pairing, "G1");
                ele_class t(pairing, "G1");
                t.set_str(xtagi[j].get<string>());
                for (auto it = ownerJson["XSet"].begin(); it != ownerJson["XSet"].end(); ++it) {
                    
                    x.set_str(it->get<string>());
                    uy = uy * (x - t);
                    fs = fs * (x + s);
                }
                wy = g ^ ((fs - uy) / (t + s));
                ownerJson["proof"].push_back(uy.to_string());
                ownerJson["proof"].push_back(wy.to_string());
            }
        }
    }
    // 这里是混的
    vector<ele_class> res;

    for (auto& item : userJson["Rw1"]) {
        ele_class t(pairing, "Zr");
        t.set_str(item.get<string>(), 16);
        res.push_back(t);
    }
    BMapAcc acc(g, s);
    ele_class lhs(acc(res));
    ele_class rhs(pairing, "G2");
    size_t n = ownerJson["proof"].size() / 2;
    n += 3;
    ownerJson["res"].clear();
    while (--n > 0) {
        rhs.random();
        ownerJson["res"].push_back(pairing(lhs, rhs).to_string());
    }

}
